Depends on firedrakeproject/fiat#262.

Summary

• gem.JaggedIndex represents simplex lattice indices with bounds depending on parent indices. TSFC lowers these to parameterized Loopy domains.
• gem.FlattenedTensor exposes a jagged lattice through a flat DoF index while retaining its tensor-product structure.
• gem.sparse_matrix represents sparse basis recombination.
• DuffyElement.duffy_evaluation() constructs sum-factorized Legendre and IntegratedLegendre tabulations; get_sparse_coeffs() provides DG ordering or CG recombination data.
• TSFC uses Duffy tabulations for coefficient and argument translation and selects collapsed quadrature for Duffy elements.

gem.optimise now:

• rewrites flat contractions as jagged lattice contractions with unflatten() and unflatten_returns();
• exposes both argument lattices before expanding separable derivative sums in bilinear forms;
• refactorizes eligible sparse-transform expressions as monomial sums and applies the same staged COFFEE sum factorization used by TSFC;
• retains the established local-distribution path for dense or residual argument transforms;
• cancels sparse row/column deltas during contraction; and
• lowers remaining FlattenedTensors to gather tables with replace_flattened().

For a three-dimensional collapsed basis, the reference derivative has 3 + 2 + 1 one-dimensional derivative paths. A bilinear Laplacian combines two such arguments and three physical gradient components. The previous unflatten_returns() expanded the first argument before exposing the second, cloning the second jagged lattice once per derivative term. The new ordering exposes both lattices jointly, then groups equal monomials and factorizes one contraction index at a time.

The final output loops remain genuinely jagged. In particular, the innermost tetrahedral domain has a bound of the form 0 <= k <= p - i - j; it is not replaced by a rectangular loop with masking.

Bilinear Laplacian code-generation benchmarks

These are tetrahedral Bernstein kernels generated in spectral mode. Structural counters and FLOPs are exact for the generated Loopy domains; times are representative local instrumented compiler runs.

Degree 10 is the tested crossover where the sum-factorized bilinear Laplacian beats canonical quadrature in both generated FLOPs and compiler time. A regression test also checks the compact instruction/temporary/loop counts and requires a two-parent parameterized domain, which preserves the tetrahedral JaggedIndex loop.

One-element execution benchmarks

FLOPs use the exact generated jagged loop domains. Times are median assembly times after warm-up.

AI assistance: OpenAI Codex was used for code analysis, implementation support, testing, and benchmarking.